package com.yuxl.common.socket.handler;

import com.yuxl.common.core.domain.auth.JwtUser;
import com.yuxl.common.auth.utils.JwtUtil;
import com.yuxl.common.core.domain.socket.Connects;
import com.yuxl.common.socket.config.WebSocketProperties;
import com.yuxl.common.socket.dispatcher.MessageDispatcher;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketSession;
import org.springframework.web.reactive.socket.WebSocketMessage;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import java.time.Duration;

@Component
public class PaopaoWebSocketHandler implements WebSocketHandler {

    @Resource
    private JwtUtil jwtUtil;
    @Resource
    private MessageDispatcher messageDispatcher;
    @Resource
    private WebSocketProperties webSocketProperties;

    @Override
    public Mono<Void> handle(WebSocketSession session) {
        // 认证逻辑：token在url参数token=xxx
        String query = session.getHandshakeInfo().getUri().getQuery();
        String token = null;
        if (query != null) {
            for (String param : query.split("&")) {
                if (param.startsWith("token=")) {
                    token = param.substring(6);
                    break;
                }
            }
        }
        if (token == null || !jwtUtil.validateToken(token)) {
            WebSocketMessage errorMsg = session.textMessage("""
                                                {"path":"/401","data":"认证失败，token无效"}
                        """);
            return session.send(Mono.just(errorMsg)).then(session.close());
        }
        JwtUser jwt = jwtUtil.getJwtFromToken(token);
        Connects.addSession(jwt.username(), session);
        // 心跳包，每30秒发送一次
        Flux<WebSocketMessage> heartbeat = Flux.interval(Duration.ofSeconds(webSocketProperties.getHeartbeatInterval()))
                .map(tick -> session.pingMessage(bufferFactory ->
                        bufferFactory.wrap(new byte[]{0x01})));
                        
                        
        // 处理消息（副流，onErrorResume防止异常导致主流终止）
        Flux<WebSocketMessage> output = session.receive()
                .filter(msg ->
                    !msg.getType().equals(WebSocketMessage.Type.PONG)
                )
                .flatMap(msg -> messageDispatcher.dispatch(jwt, session, msg))
                .onErrorResume(e -> Mono.empty());
        // 以心跳为主流，消息为副流，保证连接不断
        Flux<WebSocketMessage> merged = Flux.merge(heartbeat, output);
        return session.send(merged)
                .doFinally(s -> Connects.removeSession(jwt.username(), session));
    }
}